﻿#include <iostream>
#include <opencv2/opencv.hpp>

// 判断颜色是否为红色
bool isRedColor(cv::Vec3b color) {
    int b = color[0];
    int g = color[1];
    int r = color[2];
    return r > 150 && g < 100 && b < 100;
}

// 判断颜色是否为蓝色
bool isBlueColor(cv::Vec3b color) {
    int b = color[0];
    int g = color[1];
    int r = color[2];
    return r < 100 && g < 100 && b > 150;
}

// 灯条识别函数
void detectLightBars(cv::Mat frame,const std::string& color1) {
   
    cv::Mat hsvFrame;
    cv::cvtColor(frame, hsvFrame, cv::COLOR_BGR2HSV);

    //灰度图转化，边缘检测只能在灰度图中进行
    cv::Mat grayFrame;
    cv::cvtColor(frame, grayFrame, cv::COLOR_BGR2GRAY);

    //二值化处理
    cv::Mat binaryFrame;
    //cv::threshold(grayFrame, binaryFrame, 127, 255, cv::THRESH_BINARY);
    cv::adaptiveThreshold(grayFrame, binaryFrame, 255, cv::ADAPTIVE_THRESH_GAUSSIAN_C, cv::THRESH_BINARY, 11, 2);

    //创建结构元素用于膨胀
    cv::Mat structuringElement = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5));

    //进行膨胀操作
    cv::dilate(binaryFrame, binaryFrame, structuringElement);
    
    //高斯模糊
    cv::Mat blurredFrame;
    cv::GaussianBlur(grayFrame, blurredFrame, cv::Size(5, 5), 0);

    //检测图像边缘
    cv::Mat edges;
    cv::Canny(blurredFrame, edges, 50, 150);

    //找到边缘连接轮廓,定义轮廓列表,每个轮廓由一系列点组成
    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(edges, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    std::vector<cv::Rect> redLightBars;//定义红色灯条矩形集合
    std::vector<cv::Rect> blueLightBars;//定义蓝色灯条矩形集合

    for (const auto& contour : contours) { //遍历所有轮廓
        cv::Rect boundingRect = cv::boundingRect(contour); //计算轮廓的边界矩形
        float aspectRatio = (float)boundingRect.width / boundingRect.height; //计算宽高比
        double area = cv::contourArea(contour);
        if (area > 80 && aspectRatio < 0.4) { //如果宽高比小于0.5,则认为是灯条
            // 检查灯条颜色 //左上角坐标x,y,宽度高度
            for (int y = boundingRect.y; y < boundingRect.y + boundingRect.height; y++) { //遍历灯条区域内像素
                for (int x = boundingRect.x; x < boundingRect.x + boundingRect.width; x++) { //遍历灯条区域内像素
                    if (x < frame.cols && y < frame.rows) { //确保图像坐标在图像范围内
                        cv::Vec3b color = frame.at<cv::Vec3b>(y, x); //获取像素颜色
                        if (color1 == "red" && isRedColor(color)) { //如果为红色
                            redLightBars.push_back(boundingRect); //添加到红色灯条列表
                            break; //跳出内层循环
                        }
                        else if (color1 == "blue" && isBlueColor(color)) { //如果为蓝色
                            blueLightBars.push_back(boundingRect);  //添加到蓝色灯条列表
                            break; //跳出内层循环
                        }
                    }
                }
            }
        }
    }

    // 绘制红色灯条
    for (const auto& lightBar : redLightBars) { //遍历灯条红色列表
        cv::rectangle(frame, lightBar, cv::Scalar(0, 255, 0), 2); //在图像上绘制红色矩形框
    }

    // 绘制蓝色灯条
    for (const auto& lightBar : blueLightBars) { //遍历蓝色灯条列表
        cv::rectangle(frame, lightBar, cv::Scalar(0, 255, 0), 2); //在图像上绘制蓝色矩形边框
    }

    cv::imshow("Light Bars", frame);
}

int main() {
    std::string videoPath = "C:\\Users\\aaa\\Pictures\\Screenshots\\步兵素材红车旋转-ev--3.MOV";
    //std::string videoPath = "C:\\Users\\aaa\\Pictures\\Screenshots\\无人机航拍素材2-ev--3.MOV";
    //std::string videoPath = "C:\\Users\\aaa\\Pictures\\Screenshots\\步兵素材蓝车旋转-ev-+3.MOV";
    cv::VideoCapture video(videoPath);

    if (!video.isOpened()) {
        std::cout << "Could not open the video." << std::endl;
        return -1;
    }

    cv::Mat frame;
    std::string enermycolor = "red";
    while (true) {
        video >> frame;
        if (frame.empty())
            break;

        detectLightBars(frame, enermycolor);

        if (cv::waitKey(30) >= 0)
            break;
    }

    video.release();
    cv::destroyAllWindows();

    return 0;
}